import time

"""模型配置参数"""


class Params(object):
    """参数配置"""

    def __init__(self, model_name, device):
        """初始化"""
        """model_name模型名称(类型)"""
        self.model_name = model_name

        """使用设备(gpu设置)"""
        self.device = device

        """模型参数设置"""
        self.in_dim_node = 2  # 图中节点的维度
        self.in_dim_edge = 1  # 图中边的维度

        # 保证各个模型的参数量基本一致
        if model_name == "GCN":
            self.hidden_dim = 120  # 隐藏层维度
            self.out_dim = 120  # 输出维度
        elif model_name == "GAT":
            self.hidden_dim = 15  # 隐藏层维度(GAT使用)
            self.out_dim = 120  # 输出维度
            self.n_heads = 8  # 注意力头的个数(GAT使用)
        elif model_name == "GatedGCN":
            self.hidden_dim = 65  # 隐藏层维度
            self.out_dim = 65  # 输出维度
        else:
            raise ValueError("Not Have This Model !")

        self.n_layers = 4  # 卷积层数
        self.dropout = 0.0  # 卷积层中的dropout参数
        self.in_feat_dropout = 0.1  # 嵌入层dropout参数
        self.residual = True  # 是否使用残差连接
        self.batch_norm = True  # 是否使用批归一化

        self.device = device  # 设备配置

        """训练相关设置"""
        self.seed = 123  # 随机种子, 保证每次训练结果相同
        self.init_lr = 0.001  # 初始学习率
        self.weight_decay = 0.0001  # weight_decay防止过拟合

        self.num_epoch = 100  # epoch数(large:500)
        self.batch_size = 16  # 批大小(large:64)
        self.n_classes = 2  # 分类类别数

        # 是否使用变化lr，(训练过程中lr下降)
        self.lr_decay = False  # (large:True)
        self.lr_reduce_factor = 0.5
        self.lr_schedule_patience = 10
        self.min_lr = 1e-5

        """保存模型参数设置"""
        self.save_model = True
        self.model_path = "./checkpoints/" + model_name + "/"

        """log文件所在位置"""
        self.need_log = True
        cur_time = time.strftime("%Y-%m-%d_%H-%M-%S", time.localtime())
        self.log_file = "./logs/" + model_name + "/" + cur_time + ".log"

        """数据所在位置设置"""
        # 训练数据所在位置
        self.train_path = "./dataset/small/train_data.json"
        # 验证数据所在位置
        self.val_path = "./dataset/small/val_data.json"
        # 测试数据所在位置
        self.test_path = "./dataset/small/test_data.json"
        if self.train_path == self.val_path:
            print("train_file == val_file !")
        if self.val_path == self.test_path:
            print("val_file == test_file !")

    def get_model_info(self, num_params=None):
        """获取模型信息"""
        model_info = "\n"
        model_info += "-" * 42
        model_info += "\n模型名称: " + self.model_name
        model_info += "\n参数数量: " + str(num_params)

        model_info += "\nnum_epoch: " + str(self.num_epoch) + "\tbatch_size: " + str(self.batch_size)
        model_info += "\ninit_lr: " + str(self.init_lr) + "\tweight_decay: " + str(self.weight_decay)
        model_info += "\n"
        model_info += "-" * 42

        model_info += "\n输入节点维度: " + str(self.in_dim_node) + \
                      "\t输入边维度: " + str(self.in_dim_edge) + \
                      "\t分类个数: " + str(self.n_classes)

        model_info += "\n隐藏层维度: " + str(self.hidden_dim) + \
                      "\t输出层维度: " + str(self.out_dim) + \
                      "\t网络层数: " + str(self.n_layers)

        if self.model_name == "GAT":
            model_info += "\nnum_heads: " + str(self.n_heads)

        model_info += "\nin_feat_dropout: " + str(self.in_feat_dropout) + \
                      "\tdropout: " + str(self.dropout)
        model_info += "\nbatch_norm: " + str(self.batch_norm) + \
                      "\tresidual: " + str(self.residual)
        model_info += "\n"
        model_info += "-" * 42

        return model_info
